/*
 * Copyright 2014 The Closure Compiler Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.javascript.jscomp;

import static com.google.common.base.Preconditions.checkState;
import static com.google.common.truth.Truth.assertThat;
import static com.google.javascript.rhino.testing.NodeSubject.assertNode;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.javascript.jscomp.CompilerOptions.LanguageMode;
import com.google.javascript.jscomp.colors.Color;
import com.google.javascript.jscomp.colors.StandardColors;
import com.google.javascript.jscomp.testing.CodeSubTree;
import com.google.javascript.jscomp.testing.TestExternsBuilder;
import com.google.javascript.rhino.Node;
import com.google.javascript.rhino.testing.NodeSubject;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Unit tests for {@link Es6RewriteGenerators}. */
@RunWith(JUnit4.class)
public final class Es6RewriteGeneratorsTest extends CompilerTestCase {

  /**
   * Map from special variable names to placeholders.
   *
   * <p>`Es6RewriteGenerators` has to generate unique variable names using
   * `compiler.getUniqueIdSupplier().getUniqueId(compilerInput)`. In order to test this, we will use
   * a utility function that replaces special strings in the expected code with a combination of a
   * prefix and a compiler-input-specific hash.
   *
   * <p>This map defines the special strings that should appear in the test cases. Note that the
   * actual variables generated by `Es6RewriteGenerators` will include an additional "$" + INDEX
   * string on the end. Unless there's more than one generator function in a test case, INDEX will
   * be "0", so you would refer to the context value in the expected output code as "GEN_CONTEXT$0".
   * Each generator function rewritten within a single test case will get a unique INDEX in the
   * order that they are visited.
   */
  private static final ImmutableMap<String, String> SPECIAL_VARIABLES_MAP =
      ImmutableMap.of(
          "GEN_ARGUMENTS", "$jscomp$generator$arguments$",
          "GEN_CONTEXT", "$jscomp$generator$context$",
          "GEN_FORIN", "$jscomp$generator$forin$",
          "GEN_FUNC", "$jscomp$generator$function$",
          "GEN_THIS", "$jscomp$generator$this$");

  public Es6RewriteGeneratorsTest() {
    super(
        new TestExternsBuilder()
            .addAsyncIterable()
            .addArray()
            .addArguments()
            .addObject()
            .addMath()
            .addExtra(
                // stubs of runtime libraries
                lines(
                    "/** @const */",
                    "var $jscomp = {};",
                    "$jscomp.generator = {};",
                    "$jscomp.generator.createGenerator = function() {};",
                    "/** @constructor */",
                    "$jscomp.generator.Context = function() {};",
                    "/** @constructor */",
                    "$jscomp.generator.Context.PropertyIterator = function() {};",
                    "$jscomp.asyncExecutePromiseGeneratorFunction = function(program) {};"))
            .build());
  }

  @Override
  @Before
  public void setUp() throws Exception {
    super.setUp();
    setAcceptedLanguage(LanguageMode.ECMASCRIPT_2015);
    enableNormalize();
    enableTypeCheck();
    enableTypeInfoValidation();
    replaceTypesWithColors();
    enableMultistageCompilation();
  }

  @Override
  protected CompilerOptions getOptions() {
    CompilerOptions options = super.getOptions();
    options.setLanguageOut(LanguageMode.ECMASCRIPT3);
    return options;
  }

  @Override
  protected CompilerPass getProcessor(final Compiler compiler) {
    return new Es6RewriteGenerators(compiler);
  }

  private void rewriteGeneratorBody(String beforeBody, String afterBody) {
    rewriteGeneratorBodyWithVars(beforeBody, "", afterBody);
  }

  /** Verifies that generator functions are rewritten to a state machine program. */
  private void rewriteGeneratorBodyWithVars(String beforeBody, String varDecls, String afterBody) {
    rewriteGeneratorBodyWithVarsAndReturnType(beforeBody, varDecls, afterBody, "?");
  }

  private void rewriteGeneratorBodyWithVarsAndReturnType(
      String beforeBody, String varDecls, String afterBody, String returnType) {
    Sources srcs =
        srcs(
            lines(
                "/** @return {" + returnType + "} */",
                "function *genTestFunction() {",
                beforeBody,
                "}"));
    Expected originalExpected =
        expected(
            lines(
                "function genTestFunction() {",
                varDecls,
                "  return $jscomp.generator.createGenerator(",
                "      genTestFunction,",
                "      function (GEN_CONTEXT$0) {",
                afterBody,
                "      });",
                "}"));
    rewriteGeneratorsTest(srcs, originalExpected);
  }

  private void rewriteGeneratorsTest(Sources srcs, Expected originalExpected) {
    Expected modifiedExpected =
        expected(
            UnitTestUtils.updateGenericVarNamesInExpectedFiles(
                (FlatSources) srcs, originalExpected, SPECIAL_VARIABLES_MAP));
    test(srcs, modifiedExpected);
  }

  private void rewriteGeneratorSwitchBody(String beforeBody, String afterBody) {
    rewriteGeneratorSwitchBodyWithVars(beforeBody, "", afterBody);
  }

  /**
   * Verifies that generator functions are rewriteen to a state machine program contining {@code
   * switch} statement.
   *
   * <p>This is the case when total number of program states is more than 3.
   */
  private void rewriteGeneratorSwitchBodyWithVars(
      String beforeBody, String varDecls, String afterBody) {
    rewriteGeneratorBodyWithVars(
        beforeBody,
        varDecls,
        lines("switch (GEN_CONTEXT$0.nextAddress) {", "  case 1:", afterBody, "}"));
  }

  @Test
  public void testGeneratorForAsyncFunction() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "f = function() {",
                "  return $jscomp.asyncExecutePromiseGeneratorFunction(",
                "      function *() {",
                "        var x = 6;",
                "        yield x;",
                "      });",
                "}")),
        expected(
            lines(
                "f = function () {",
                "  var x;",
                "  return $jscomp.asyncExecutePromiseGeneratorProgram(",
                "      function (GEN_CONTEXT$0) {",
                "         x = 6;",
                "         return GEN_CONTEXT$0.yield(x, 0);",
                "      });",
                "}")));

    rewriteGeneratorsTest(
        srcs(
            lines(
                "f = function(x) {",
                "  var $jscomp$restParams = [];",
                "  for (var $jscomp$restIndex = 0;",
                "      $jscomp$restIndex < arguments.length;",
                "      ++$jscomp$restIndex) {",
                "    $jscomp$restParams[$jscomp$restIndex - 0] = arguments[$jscomp$restIndex];",
                "  }",
                "  {",
                "    var bla$0 = $jscomp$restParams;",
                "    return $jscomp.asyncExecutePromiseGeneratorFunction(",
                "        function *() {",
                "          var y = bla$0[0];",
                "          yield y;",
                "        });",
                "  }",
                "}")),
        expected(
            lines(
                "f = function (x) {",
                "  var $jscomp$restParams = [];",
                "  var $jscomp$restIndex = 0;",
                "  for (;",
                "      $jscomp$restIndex < arguments.length;",
                "      ++$jscomp$restIndex) {",
                "    $jscomp$restParams[$jscomp$restIndex - 0] = arguments[$jscomp$restIndex];",
                "  }",
                "  {",
                "    var bla$0 = $jscomp$restParams;",
                "    var y;",
                "    return $jscomp.asyncExecutePromiseGeneratorProgram(",
                "        function (GEN_CONTEXT$0) {",
                "          y = bla$0[0];",
                "          return GEN_CONTEXT$0.yield(y, 0);",
                "        });",
                "  }",
                "}")));
  }

  @Test
  public void testUnnamed() {
    rewriteGeneratorsTest(
        srcs(lines("f = function *() {};")),
        expected(
            lines(
                "f = function GEN_FUNC$0() {",
                "  return $jscomp.generator.createGenerator(",
                "      GEN_FUNC$0,",
                "      function (GEN_CONTEXT$0) {",
                "         GEN_CONTEXT$0.jumpToEnd();",
                "      });",
                "}")));
  }

  @Test
  public void testConst() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "/** @const */", //
                "var f = function *() {};")),
        expected(
            lines(
                "/** @const */", //
                "var f = function GEN_FUNC$0() {",
                "  return $jscomp.generator.createGenerator(",
                "      GEN_FUNC$0,",
                "      function (GEN_CONTEXT$0) {",
                "         GEN_CONTEXT$0.jumpToEnd();",
                "      });",
                "}")));
  }

  @Test
  public void testContainsClosure() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "function *f(o) {", //
                "  var i = 0",
                "  var pp;",
                "  function msg(p) {",
                "    return i++ + ' ' + p + '\\n';",
                "  }",
                "  for (pp in obj) {",
                "    yield msg(pp);",
                "  }",
                "}")),
        expected(
            lines(
                "function f(o) {",
                "  function msg(p) {",
                "    return i++ + ' ' + p + '\\n';",
                "  }",
                "  var i",
                "  var pp;",
                "  var GEN_FORIN$0$0;",
                "  return $jscomp.generator.createGenerator(",
                "      f,",
                "      function (GEN_CONTEXT$0) {",
                "         if (GEN_CONTEXT$0.nextAddress == 1) {",
                "           i = 0;",
                "           GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(obj);",
                "         }",
                "         if (!((pp = GEN_FORIN$0$0.getNext()) != null)) {",
                "           return GEN_CONTEXT$0.jumpTo(0);",
                "         }",
                "         return GEN_CONTEXT$0.yield(msg(pp), 2);",
                "      });",
                "}")));
  }

  @Test
  public void testContainsClosureAssignedToVariable() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "function *f(o) {", //
                "  var i = 0",
                "  var pp;",
                "  var msg = function(p) {",
                "    return i++ + ' ' + p + '\\n';",
                "  }",
                "  for (pp in obj) {",
                "    yield msg(pp);",
                "  }",
                "}")),
        expected(
            lines(
                "function f(o) {",
                "  var i",
                "  var pp;",
                "  var msg;",
                "  var GEN_FORIN$0$0;",
                "  return $jscomp.generator.createGenerator(",
                "      f,",
                "      function (GEN_CONTEXT$0) {",
                "         if (GEN_CONTEXT$0.nextAddress == 1) {",
                "           i = 0;",
                // TODO(bradfordcsmith): Maybe it would be better if we hoisted this function
                // expression out so it doesn't create a new closure every time we enter this
                // callback function?
                "           msg = function(p) {",
                "             return i++ + ' ' + p + '\\n';",
                "           };",
                "           GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(obj);",
                "         }",
                "         if (!((pp = GEN_FORIN$0$0.getNext()) != null)) {",
                "           return GEN_CONTEXT$0.jumpTo(0);",
                "         }",
                "         return GEN_CONTEXT$0.yield(msg(pp), 2);",
                "      });",
                "}")));
  }

  @Test
  public void testContainsClosureAndTranspiledFromAsync() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "function f(o) {",
                "  use(o);",
                "  return $jscomp.asyncExecutePromiseGeneratorFunction(",
                "      function *() {", //
                "        var i = 0",
                "        var pp;",
                "        function msg(p) {",
                "          return i++ + ' ' + p + '\\n';",
                "        }",
                "        for (pp in obj) {",
                "          yield msg(pp);",
                "        }",
                "      });",
                "}")),
        expected(
            lines(
                "function f(o) {",
                // When Es6RewriteGenerators recognizes that it is transpiling code that was
                // originally an async function, it will use the body of that function as the
                // place to move variable declarations instead of putting them in the generator
                // function body itself.
                //
                // To maintain normalization, all function declarations must go to the top
                "  function msg(p) {",
                "    return i++ + ' ' + p + '\\n';",
                "  }",
                // In a real situation the only thing that is likely to be here is code that
                // was generated by transpilations that occurred after async function transpilation
                // and before generator transpilation. For example code for handling parameter
                // destructuring.
                "  use(o);",
                // Regular var declarations just go in the order they originally appeared
                // immediately before the executor call.
                "  var i",
                "  var pp;",
                "  var GEN_FORIN$0$0;",
                "  return $jscomp.asyncExecutePromiseGeneratorProgram(",
                "      function (GEN_CONTEXT$0) {",
                "         if (GEN_CONTEXT$0.nextAddress == 1) {",
                "           i = 0;",
                "           GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(obj);",
                "         }",
                "         if (!((pp = GEN_FORIN$0$0.getNext()) != null)) {",
                "           return GEN_CONTEXT$0.jumpTo(0);",
                "         }",
                "         return GEN_CONTEXT$0.yield(msg(pp), 2);",
                "      });",
                "}")));
  }

  @Test
  public void testSimpleGenerator() {
    rewriteGeneratorBody("", "  GEN_CONTEXT$0.jumpToEnd();");

    rewriteGeneratorBody("yield;", lines("  return GEN_CONTEXT$0.yield(void 0, 0);"));

    rewriteGeneratorBody("yield 1;", lines("  return GEN_CONTEXT$0.yield(1, 0);"));

    Sources srcs = srcs("/** @param {*} x */ function *f(x, y) {}");
    Expected originalExpected =
        expected(
            lines(
                "function f(x, y) {",
                "  return $jscomp.generator.createGenerator(",
                "      f,",
                "      function(GEN_CONTEXT$0) {",
                "        GEN_CONTEXT$0.jumpToEnd();",
                "      });",
                "}"));
    rewriteGeneratorsTest(srcs, originalExpected);

    rewriteGeneratorBodyWithVars(
        "var i = 0, j = 2;",
        lines(
            "var i;", //
            "var j;"),
        lines(
            "i = 0;", //
            "j = 2;",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "var i = 0; yield i; i = 1; yield i; i = i + 1; yield i;",
        "var i;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  i = 0;",
            "  return GEN_CONTEXT$0.yield(i, 2);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  i = 1;",
            "  return GEN_CONTEXT$0.yield(i, 3);",
            "}",
            "i = i + 1;",
            "return GEN_CONTEXT$0.yield(i, 0);"));
  }

  @Test
  public void testForLoopWithExtraVarDeclaration() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "function *gen() {",
                "  var i = 2;",
                "  yield i;",
                "  use(i);",
                "  for (var i = 0; i < 3; i++) {",
                "    use(i);",
                "  }",
                "}")),
        expected(
            lines(
                "function gen(){",
                "  var i;",
                "  return $jscomp.generator.createGenerator(",
                "      gen,",
                "      function(GEN_CONTEXT$0) {",
                "        if (GEN_CONTEXT$0.nextAddress==1) {",
                "          i=2;",
                "          return GEN_CONTEXT$0.yield(i,2);",
                "        }",
                "        use(i);",
                "        i = 0;",
                "        for (; i < 3 ; i++) use(i);",
                "        GEN_CONTEXT$0.jumpToEnd();",
                "      })",
                "}")));
  }

  @Test
  public void testUnreachableCodeGeneration() {
    rewriteGeneratorBody(
        "if (i) return 1; else return 2;",
        lines(
            "  if (i) {",
            "    return GEN_CONTEXT$0.return(1);",
            "  } else {",
            "    return GEN_CONTEXT$0.return(2);",
            "  }",
            // TODO(b/73762053): Avoid generating unreachable statements.
            "  GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testReturnGenerator() {
    rewriteGeneratorsTest(
        srcs("function f() { return function *g() {yield 1;} }"),
        expected(
            lines(
                "function f() {",
                "  return function g() {",
                "    return $jscomp.generator.createGenerator(",
                "        g,",
                "        function(GEN_CONTEXT$0) {",
                "          return GEN_CONTEXT$0.yield(1, 0);",
                "        });",
                "  }",
                "}")));
  }

  @Test
  public void testNestedGenerator() {
    rewriteGeneratorsTest(
        srcs("function *f() { function *g() {yield 2;} yield 1; }"),
        expected(
            lines(
                "function f() {",
                "  function g() {",
                "    return $jscomp.generator.createGenerator(",
                "        g,",
                "        function(GEN_CONTEXT$0) {",
                "          return GEN_CONTEXT$0.yield(2, 0);",
                "        });",
                "  }",
                "  return $jscomp.generator.createGenerator(",
                "      f,",
                "      function(GEN_CONTEXT$1) {",
                "        return GEN_CONTEXT$1.yield(1, 0);",
                "      });",
                "}")));
  }

  @Test
  public void testForLoops() {
    rewriteGeneratorBodyWithVars(
        "var i = 0; for (var j = 0; j < 10; j++) { i += j; }",
        "var i; var j;",
        lines(
            "i = 0;", //
            "j = 0;",
            "for (; j < 10; j++) {",
            "  i = i + j;", // normalization rewrote this
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "var i = 0; for (var j = yield; j < 10; j++) { i += j; }",
        "var i; var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  i = 0;",
            "  return GEN_CONTEXT$0.yield(void 0, 2);",
            "}",
            "j = GEN_CONTEXT$0.yieldResult;",
            "for (; j < 10; j++) {",
            "  i = i + j;", // normalization rewrote this
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBody("for (;;) { yield 1; }", lines("  return GEN_CONTEXT$0.yield(1, 1);"));

    rewriteGeneratorBodyWithVars(
        "for (var yieldResult; yieldResult === undefined; yieldResult = yield 1) {}",
        "var yieldResult;",
        lines(
            "  if (GEN_CONTEXT$0.nextAddress == 1) {",
            "    if (!(yieldResult === undefined)) return GEN_CONTEXT$0.jumpTo(0);",
            "    return GEN_CONTEXT$0.yield(1,5);",
            "  }",
            "  yieldResult = GEN_CONTEXT$0.yieldResult;",
            "  return GEN_CONTEXT$0.jumpTo(1);"));

    rewriteGeneratorBodyWithVars(
        "for (var j = 0; j < 10; j++) { yield j; }",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  if (!(j < 10)) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 3);",
            "}",
            "j++;",
            "return GEN_CONTEXT$0.jumpTo(2);"));

    rewriteGeneratorBodyWithVars(
        "var i = 0; for (var j = 0; j < 10; j++) { i += j; yield 5; }",
        "var i; var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  i = 0;",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  if (!(j < 10)) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  i = i + j;", // normalization rewrote this
            "  return GEN_CONTEXT$0.yield(5, 3);",
            "}",
            "j++;",
            "return GEN_CONTEXT$0.jumpTo(2);"));
  }

  @Test
  public void testWhileLoops() {
    rewriteGeneratorBodyWithVars(
        "var i = 0; while (i < 10) { i++; i++; i++; } yield i;",
        "  var i;",
        lines(
            "i = 0;", //
            "for (; i < 10;) { i ++; i++; i++; }",
            "return GEN_CONTEXT$0.yield(i, 0);"));

    rewriteGeneratorSwitchBodyWithVars(
        "var j = 0; while (j < 10) { yield j; j++; } j += 10;",
        "var j;",
        lines(
            "  j = 0;",
            "case 2:",
            "  if (!(j < 10)) {",
            "    GEN_CONTEXT$0.jumpTo(4);",
            "    break;",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 5)",
            "case 5:",
            "  j++;",
            "  GEN_CONTEXT$0.jumpTo(2);",
            "  break;",
            "case 4:",
            "  j = j + 10;", // normalization rewrote this
            "  GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "var j = 0; while (j < 10) { yield j; j++; } yield 5",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 5) {",
            "  if (!(j < 10)) {",
            "    return GEN_CONTEXT$0.yield(5, 0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 5)",
            "}",
            "  j++;",
            "  return GEN_CONTEXT$0.jumpTo(2);"));

    rewriteGeneratorBodyWithVars(
        "var j = 0; while (yield) { j++; }",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 5) {",
            "  return GEN_CONTEXT$0.yield(void 0, 5);",
            "}",
            "if (!(GEN_CONTEXT$0.yieldResult)) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "j++;",
            "return GEN_CONTEXT$0.jumpTo(2);"));
  }

  @Test
  public void testDecomposeComplexYieldExpression() {
    rewriteGeneratorsTest(
        srcs(
            lines(
                "/* @return {?} */ function *f() {",
                "  var o = {bar: function(x) {}};",
                "  (yield 5) && o.bar(yield 5);",
                "}")),
        expected(
            lines(
                "function f(){",
                "  var o;",
                "  var JSCompiler_temp$jscomp$0;",
                "  var JSCompiler_temp_const$jscomp$2;",
                "  var JSCompiler_temp_const$jscomp$1;",
                "  var JSCompiler_temp_const$jscomp$3;",
                "  return $jscomp.generator.createGenerator(f, function(GEN_CONTEXT$0) {",
                "    switch(GEN_CONTEXT$0.nextAddress) {",
                "      case 1:",
                "        o = {bar:function(x) {}};",
                "        return GEN_CONTEXT$0.yield(5, 2);",
                "      case 2:",
                "        if (!(JSCompiler_temp$jscomp$0 = GEN_CONTEXT$0.yieldResult)) {",
                "          GEN_CONTEXT$0.jumpTo(3);",
                "          break;",
                "        }",
                "        JSCompiler_temp_const$jscomp$2 = o;",
                "        JSCompiler_temp_const$jscomp$1 = JSCompiler_temp_const$jscomp$2.bar;",
                "        JSCompiler_temp_const$jscomp$3 = JSCompiler_temp_const$jscomp$2;",
                "        return GEN_CONTEXT$0.yield(5, 4);",
                "      case 4:",
                "        JSCompiler_temp$jscomp$0 =",
                "            JSCompiler_temp_const$jscomp$1.call(",
                "                JSCompiler_temp_const$jscomp$3,",
                "                GEN_CONTEXT$0.yieldResult);",
                "      case 3:",
                "        JSCompiler_temp$jscomp$0;",
                "        GEN_CONTEXT$0.jumpToEnd();",
                "    }",
                "  });",
                "}")));
  }

  @Test
  public void testDecomposableExpression() {
    rewriteGeneratorBodyWithVars(
        "return a + (a = b) + (b = yield) + a;",
        lines("var JSCompiler_temp_const$jscomp$0;"),
        lines(
            "  if (GEN_CONTEXT$0.nextAddress == 1) {",
            "    JSCompiler_temp_const$jscomp$0 = a + (a = b);",
            "    return GEN_CONTEXT$0.yield(void 0, 2);",
            "  }",
            "  return GEN_CONTEXT$0.return(",
            "JSCompiler_temp_const$jscomp$0 + (b = GEN_CONTEXT$0.yieldResult) + a);"));

    rewriteGeneratorSwitchBodyWithVars(
        "return (yield ((yield 1) + (yield 2)));",
        lines("var JSCompiler_temp_const$jscomp$0;"),
        lines(
            "  return GEN_CONTEXT$0.yield(1, 3);",
            "case 3:",
            "  JSCompiler_temp_const$jscomp$0=GEN_CONTEXT$0.yieldResult;",
            "  return GEN_CONTEXT$0.yield(2, 4);",
            "case 4:",
            "  return GEN_CONTEXT$0.yield(",
            "      JSCompiler_temp_const$jscomp$0 + GEN_CONTEXT$0.yieldResult, 2);",
            "case 2:",
            "  return GEN_CONTEXT$0.return(GEN_CONTEXT$0.yieldResult);"));

    rewriteGeneratorBodyWithVars(
        "var o = {bar: function(x) {}}; o.bar(yield 5);",
        lines(
            "var o;", "var JSCompiler_temp_const$jscomp$1;", "var JSCompiler_temp_const$jscomp$0;"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  o = {bar: function(x) {}};",
            "  JSCompiler_temp_const$jscomp$1 = o;",
            "  JSCompiler_temp_const$jscomp$0 = JSCompiler_temp_const$jscomp$1.bar;",
            "  return GEN_CONTEXT$0.yield(5, 2);",
            "}",
            "JSCompiler_temp_const$jscomp$0.call(",
            "    JSCompiler_temp_const$jscomp$1, GEN_CONTEXT$0.yieldResult);",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testGeneratorCannotConvertYet() {
    testError(
        "function *f(b, i) {switch (i) { case yield: return b; }}",
        TranspilationUtil.CANNOT_CONVERT_YET);
  }

  @Test
  public void testThrow() {
    rewriteGeneratorBody("throw 1;", "throw 1;");
  }

  @Test
  public void testLabels() {
    rewriteGeneratorBody(
        "l: if (true) { break l; }",
        lines(
            "  l: {", //
            "    if (true) { break l; }",
            "  }",
            "  GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBody(
        "l: if (yield) { break l; }",
        lines(
            "  if (GEN_CONTEXT$0.nextAddress == 1)",
            "    return GEN_CONTEXT$0.yield(void 0, 3);",
            "  if (GEN_CONTEXT$0.yieldResult) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBody(
        "l: if (yield) { while (1) {break l;} }",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(void 0, 3);",
            "}",
            "if (GEN_CONTEXT$0.yieldResult) {",
            "  for (; 1;) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBody(
        "l: for (;;) { yield i; continue l; }",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(i, 5);",
            "}",
            "return GEN_CONTEXT$0.jumpTo(1);",
            "return GEN_CONTEXT$0.jumpTo(1);"));

    rewriteGeneratorBody(
        "l1: l2: if (yield) break l1; else break l2;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(void 0, 3);",
            "}",
            "if(GEN_CONTEXT$0.yieldResult) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "} else {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testUnreachable() {
    // TODO(skill): The generator transpilation should not produce any unreachable code
    rewriteGeneratorBody(
        "while (true) {yield; break;}",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  if (!true) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(void 0, 5);",
            "}",
            "return GEN_CONTEXT$0.jumpTo(0);",
            "return GEN_CONTEXT$0.jumpTo(1);"));
  }

  @Test
  public void testCaseNumberOptimization() {
    rewriteGeneratorBodyWithVars(
        lines(
            "for (;;) {",
            "  var gen = generatorSrc();",
            "  var gotResponse = false;",
            "  for (var response in gen) {",
            "    yield response;",
            "    gotResponse = true;",
            "  }",
            "  if (!gotResponse) {",
            "    return;",
            "  }",
            "}"),
        lines(
            "var gen;", //
            "var gotResponse;",
            "var response;",
            "var GEN_FORIN$0$0;"),
        lines(
            "switch (GEN_CONTEXT$0.nextAddress) {",
            "  case 1:",
            "    gen = generatorSrc();",
            "    gotResponse = false;",
            "    GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(gen);",
            "  case 5:",
            "    if (!((response=GEN_FORIN$0$0.getNext()) != null)) {",
            "      GEN_CONTEXT$0.jumpTo(7);",
            "      break;",
            "    }",
            "    return GEN_CONTEXT$0.yield(response, 8);",
            "  case 8:",
            "    gotResponse = true;",
            "    GEN_CONTEXT$0.jumpTo(5);",
            "    break;",
            "  case 7:",
            "    if (!gotResponse) return GEN_CONTEXT$0.return();",
            "    GEN_CONTEXT$0.jumpTo(1);",
            "    break;",
            "}",
            ""));

    rewriteGeneratorBody(
        "do { do { do { yield; } while (3) } while (2) } while (1)",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(void 0, 10);",
            "}",
            "if (3) {",
            "  return GEN_CONTEXT$0.jumpTo(1);",
            "}",
            "if (2) {",
            "  return GEN_CONTEXT$0.jumpTo(1);",
            "}",
            "if (1) {",
            "  return GEN_CONTEXT$0.jumpTo(1);",
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testIf() {
    rewriteGeneratorBodyWithVars(
        "var j = 0; if (yield) { j = 1; }",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "  return GEN_CONTEXT$0.yield(void 0, 2);",
            "}",
            "if (GEN_CONTEXT$0.yieldResult) { j = 1; }",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "var j = 0; if (j < 1) { j = 5; } else { yield j; }",
        "var j;",
        lines(
            "j = 0;",
            "if (j < 1) {",
            "  j = 5;",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(j, 0);"));

    // When "else" doesn't contain yields, it's more optimal to swap "if" and else "blocks" and
    // negate the condition.
    rewriteGeneratorBodyWithVars(
        "var j = 0; if (j < 1) { yield j; } else { j = 5; }",
        "var j;",
        lines(
            "j = 0;",
            "if (!(j < 1)) {",
            "  j = 5;",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(j, 0);"));

    // No "else" block, pretend as it's empty
    rewriteGeneratorBodyWithVars(
        "var j = 0; if (j < 1) { yield j; }",
        "var j;",
        lines(
            "j = 0;",
            "if (!(j < 1)) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(j, 0);"));

    rewriteGeneratorBody(
        "if (i < 1) { yield i; } else { yield 1; }",
        lines(
            "if (i < 1) {",
            "  return GEN_CONTEXT$0.yield(i, 0);",
            "}",
            "return GEN_CONTEXT$0.yield(1, 0);"));

    rewriteGeneratorSwitchBody(
        "if (i < 1) { yield i; yield i + 1; i = 10; } else { yield 1; yield 2; i = 5;}",
        lines(
            "  if (i < 1) {",
            "    return GEN_CONTEXT$0.yield(i, 6);",
            "  }",
            "  return GEN_CONTEXT$0.yield(1, 4);",
            "case 4:",
            "  return GEN_CONTEXT$0.yield(2, 5);",
            "case 5:",
            "  i = 5;",
            "  GEN_CONTEXT$0.jumpTo(0);",
            "  break;",
            "case 6:",
            "  return GEN_CONTEXT$0.yield(i + 1, 7)",
            "case 7:",
            "  i = 10;",
            "  GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBody(
        lines(
            "if (i < 1) {",
            "  while (true) {",
            "    yield 1;",
            "  }",
            "} else {",
            "  while (false) {",
            "    yield 2;",
            "  }",
            "}"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  if (i < 1) {",
            "    return GEN_CONTEXT$0.jumpTo(8);",
            "  }",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 8) {",
            "  if (!false) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(2, 4);",
            "}",
            "if (!true) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(1, 8);"));

    rewriteGeneratorBody(
        "if (i < 1) { if (i < 2) {yield 2; } } else { yield 1; }",
        lines(
            "if (i < 1) {",
            "  if (!(i < 2)) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(2, 0);",
            "}",
            "return GEN_CONTEXT$0.yield(1, 0)"));

    rewriteGeneratorBody(
        "if (i < 1) { if (i < 2) {yield 2; } else { yield 3; } } else { yield 1; }",
        lines(
            "if (i < 1) {",
            "  if (i < 2) {",
            "    return GEN_CONTEXT$0.yield(2, 0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(3, 0);",
            "}",
            "return GEN_CONTEXT$0.yield(1, 0)"));
  }

  @Test
  public void testReturn() {
    rewriteGeneratorBody("return 1;", "return GEN_CONTEXT$0.return(1);");

    rewriteGeneratorBodyWithVars(
        "return this;",
        "var GEN_THIS$0 = this;",
        lines("return GEN_CONTEXT$0.return(GEN_THIS$0);"));

    rewriteGeneratorBodyWithVars(
        "return this.test({value: this});",
        "var GEN_THIS$0 = this;",
        lines("return GEN_CONTEXT$0.return(", "    GEN_THIS$0.test({value: GEN_THIS$0}));"));

    rewriteGeneratorBodyWithVars(
        "return this[yield];",
        "var GEN_THIS$0 = this;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1)",
            "  return GEN_CONTEXT$0.yield(void 0, 2);",
            "return GEN_CONTEXT$0.return(",
            "    GEN_THIS$0[GEN_CONTEXT$0.yieldResult]);"));
  }

  @Test
  public void testBreakContinue() {
    rewriteGeneratorBodyWithVars(
        "var j = 0; while (j < 10) { yield j; break; }",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 5) {",
            "  if (!(j < 10)) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 5);",
            "}",
            "return GEN_CONTEXT$0.jumpTo(0);",
            "return GEN_CONTEXT$0.jumpTo(2)"));

    rewriteGeneratorBodyWithVars(
        "var j = 0; while (j < 10) { yield j; continue; }",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 5) {",
            "  if (!(j < 10)) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 5);",
            "}",
            "return GEN_CONTEXT$0.jumpTo(2);",
            "return GEN_CONTEXT$0.jumpTo(2)"));

    rewriteGeneratorBodyWithVars(
        "for (var j = 0; j < 10; j++) { yield j; break; }",
        "var j;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  j = 0;",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 5) {",
            "  if (!(j < 10)) {",
            "    return GEN_CONTEXT$0.jumpTo(0);",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 5);",
            "}",
            "return GEN_CONTEXT$0.jumpTo(0);",
            "j++;",
            "return GEN_CONTEXT$0.jumpTo(2)"));

    rewriteGeneratorSwitchBodyWithVars(
        "for (var j = 0; j < 10; j++) { yield j; continue; }",
        "var j;",
        lines(
            "  j = 0;",
            "case 2:",
            "  if (!(j < 10)) {",
            "    GEN_CONTEXT$0.jumpTo(0);",
            "    break;",
            "  }",
            "  return GEN_CONTEXT$0.yield(j, 5);",
            "case 5:",
            "  GEN_CONTEXT$0.jumpTo(3);",
            "  break;",
            "case 3:",
            "  j++;",
            "  GEN_CONTEXT$0.jumpTo(2)",
            "  break;"));
  }

  @Test
  public void testDoWhileLoops() {
    rewriteGeneratorBody(
        "do { yield j; } while (j < 10);",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(j, 4);",
            "}",
            "if (j<10) {",
            "  return GEN_CONTEXT$0.jumpTo(1);",
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBody(
        "do {} while (yield 1);",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(1, 5);",
            "}",
            "if (GEN_CONTEXT$0.yieldResult) {",
            "  return GEN_CONTEXT$0.jumpTo(1);",
            "}",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testYieldNoValue() {
    rewriteGeneratorBody("yield;", "return GEN_CONTEXT$0.yield(void 0, 0);");
  }

  @Test
  public void testReturnNoValue() {
    rewriteGeneratorBody(
        "return;", //
        "return GEN_CONTEXT$0.return();");
  }

  @Test
  public void testYieldExpression() {
    rewriteGeneratorBody(
        "return (yield 1);",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(1, 2);",
            "}",
            "return GEN_CONTEXT$0.return(GEN_CONTEXT$0.yieldResult);"));
  }

  @Test
  public void testFunctionInGenerator() {
    rewriteGeneratorBodyWithVars(
        "function g() {}", // body statement is just s function definition
        "function g() {}", // and it remains the only variable declaration after transpilation
        "  GEN_CONTEXT$0.jumpToEnd();");
  }

  @Test
  public void testYieldAll() {
    rewriteGeneratorBody(
        "yield * n;", //
        "  return GEN_CONTEXT$0.yieldAll(n, 0);");

    rewriteGeneratorBodyWithVars(
        "var i = yield * n;",
        "var i;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yieldAll(n, 2);",
            "}",
            "i = GEN_CONTEXT$0.yieldResult;",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testYieldArguments() {
    rewriteGeneratorBodyWithVars(
        "yield arguments[0];",
        "var GEN_ARGUMENTS$0 = arguments;",
        "return GEN_CONTEXT$0.yield(GEN_ARGUMENTS$0[0], 0);");
  }

  @Test
  public void testYieldThis() {
    rewriteGeneratorBodyWithVars(
        "yield this;", "var GEN_THIS$0 = this;", "return GEN_CONTEXT$0.yield(GEN_THIS$0, 0);");
  }

  @Test
  public void testGeneratorShortCircuit() {
    rewriteGeneratorBodyWithVars(
        "0 || (yield 1);",
        "var JSCompiler_temp$jscomp$0;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  if(JSCompiler_temp$jscomp$0 = 0) {",
            "    return GEN_CONTEXT$0.jumpTo(2);",
            "  }",
            "  return GEN_CONTEXT$0.yield(1, 3);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 2) {",
            "  JSCompiler_temp$jscomp$0=GEN_CONTEXT$0.yieldResult;",
            "}",
            "JSCompiler_temp$jscomp$0;",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "0 && (yield 1);",
        "var JSCompiler_temp$jscomp$0;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  if(!(JSCompiler_temp$jscomp$0=0)) {",
            "    return GEN_CONTEXT$0.jumpTo(2);",
            "  }",
            "  return GEN_CONTEXT$0.yield(1, 3);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 2) {",
            "  JSCompiler_temp$jscomp$0=GEN_CONTEXT$0.yieldResult;",
            "}",
            "JSCompiler_temp$jscomp$0;",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "0 ? 1 : (yield 1);",
        "var JSCompiler_temp$jscomp$0;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  if(0) {",
            "    JSCompiler_temp$jscomp$0 = 1;",
            "    return GEN_CONTEXT$0.jumpTo(2);",
            "  }",
            "  return GEN_CONTEXT$0.yield(1, 3);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 2) {",
            "  JSCompiler_temp$jscomp$0 = GEN_CONTEXT$0.yieldResult;",
            "}",
            "JSCompiler_temp$jscomp$0;",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testVar() {
    rewriteGeneratorBodyWithVars(
        "var va = 10, vb, vc = yield 10, vd = yield 20, vf, vg='test';",
        lines(
            "var va;", //
            "var vb;", "var vc;", "var vd;", "var vf;", "var vg;"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  va = 10;",
            "  return GEN_CONTEXT$0.yield(10, 2);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  vc = GEN_CONTEXT$0.yieldResult;",
            "  return GEN_CONTEXT$0.yield(20, 3);",
            "}",
            "vd = GEN_CONTEXT$0.yieldResult;",
            "vg = 'test';",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        lines("var /** @const */ va = 10, vb, vc = yield 10, vd = yield 20, vf, vg='test';"),
        lines(
            "var /** @const */ va;", //
            "var vb;",
            "var vc;",
            "var vd",
            "var vf;",
            "var vg;"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  va = 10;",
            "  return GEN_CONTEXT$0.yield(10, 2);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  vc = GEN_CONTEXT$0.yieldResult;",
            "  return GEN_CONTEXT$0.yield(20, 3);",
            "}",
            "vd = GEN_CONTEXT$0.yieldResult;",
            "vg = 'test';",
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testYieldSwitch() {
    rewriteGeneratorSwitchBody(
        lines(
            "while (1) {",
            "  switch (i) {",
            "    case 1:",
            "      ++i;",
            "      break;",
            "    case 2:",
            "      yield 3;",
            "      continue;",
            "    case 10:",
            "    case 3:",
            "      yield 4;",
            "    case 4:",
            "      return 1;",
            "    case 5:",
            "      return 2;",
            "    default:",
            "      yield 5;",
            "  }",
            "}"),
        lines(
            "  if (!1) {",
            "    GEN_CONTEXT$0.jumpTo(0);",
            "    break;",
            "  }",
            "  switch (i) {",
            "    case 1: ++i; break;",
            "    case 2: return GEN_CONTEXT$0.jumpTo(5);",
            "    case 10:",
            "    case 3: return GEN_CONTEXT$0.jumpTo(6);",
            "    case 4: return GEN_CONTEXT$0.jumpTo(7);",
            "    case 5: return GEN_CONTEXT$0.return(2);",
            "    default: return GEN_CONTEXT$0.jumpTo(8);",
            "  }",
            "  GEN_CONTEXT$0.jumpTo(1);",
            "  break;",
            "case 5: return GEN_CONTEXT$0.yield(3, 10);",
            "case 10:",
            "  GEN_CONTEXT$0.jumpTo(1);",
            "  break;",
            "case 6: return GEN_CONTEXT$0.yield(4, 7);",
            "case 7: return GEN_CONTEXT$0.return(1);",
            "case 8: return GEN_CONTEXT$0.yield(5, 1);"));

    rewriteGeneratorBody(
        lines("switch (yield) {", "  default:", "  case 1:", "    yield 1;}"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(void 0, 2);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  switch (GEN_CONTEXT$0.yieldResult) {",
            "    default:",
            "    case 1:",
            "      return GEN_CONTEXT$0.jumpTo(3)",
            "  }",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(1, 0);"));
  }

  @Test
  public void testNoTranslate() {
    rewriteGeneratorBody(
        "if (1) { try {} catch (e) {} throw 1; }",
        lines(
            "if (1) { try {} catch (e) {} throw 1; }", //
            "GEN_CONTEXT$0.jumpToEnd();"));
  }

  @Test
  public void testForIn() {
    rewriteGeneratorBodyWithVars(
        "for (var i in yield) { }",
        "var i;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(void 0, 2);",
            "}",
            "for (i in GEN_CONTEXT$0.yieldResult) { }",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "for (var i in j) { yield i; }",
        lines(
            "var i;", //
            "var GEN_FORIN$0$0;"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(j);",
            "}",
            "if (!((i = GEN_FORIN$0$0.getNext()) != null)) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(i, 2);"));

    rewriteGeneratorBodyWithVars(
        "for (var i in yield) { yield i; }",
        lines(
            "var i;", //
            "var GEN_FORIN$0$0;"),
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  return GEN_CONTEXT$0.yield(void 0, 2)",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  GEN_FORIN$0$0 = ",
            "      GEN_CONTEXT$0.forIn(GEN_CONTEXT$0.yieldResult);",
            "}",
            "if (!((i = GEN_FORIN$0$0.getNext()) != null)) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.yield(i, 3);"));

    rewriteGeneratorBodyWithVars(
        "for (i[yield] in j) {}",
        "var GEN_FORIN$0$0; var JSCompiler_temp_const$jscomp$0;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn(j);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 5) {",
            "  JSCompiler_temp_const$jscomp$0 = i;",
            "  return GEN_CONTEXT$0.yield(void 0, 5);",
            "}",
            "if (!((JSCompiler_temp_const$jscomp$0[GEN_CONTEXT$0.yieldResult] =",
            "    GEN_FORIN$0$0.getNext()) != null)) {",
            "  return GEN_CONTEXT$0.jumpTo(0);",
            "}",
            "return GEN_CONTEXT$0.jumpTo(2);"));
  }

  @Test
  public void testTryCatch() {
    rewriteGeneratorBodyWithVars(
        "try {yield 1;} catch (e) {}",
        "var e;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(2);",
            "  return GEN_CONTEXT$0.yield(1, 4);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 2) {",
            "  return GEN_CONTEXT$0.leaveTryBlock(0)",
            "}",
            "e = GEN_CONTEXT$0.enterCatchBlock();",
            "GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorSwitchBodyWithVars(
        lines(
            "try {yield 1;} catch (e) { use('first ' + e); }", //
            "try {yield 1;} catch (e) { use('second ' + e)}"),
        lines(
            "var e;",
            // The second catch variable was renamed by normalization
            "var e$jscomp$1;"),
        lines(
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(2);",
            "  return GEN_CONTEXT$0.yield(1, 4);",
            "case 4:",
            "  GEN_CONTEXT$0.leaveTryBlock(3)",
            "  break;",
            "case 2:",
            "  e=GEN_CONTEXT$0.enterCatchBlock();",
            "  use('first ' + e);",
            "case 3:",
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(5);",
            "  return GEN_CONTEXT$0.yield(1, 7);",
            "case 7:",
            "  GEN_CONTEXT$0.leaveTryBlock(0)",
            "  break;",
            "case 5:",
            "  e$jscomp$1 = GEN_CONTEXT$0.enterCatchBlock();",
            "  use('second ' + e$jscomp$1);",
            "  GEN_CONTEXT$0.jumpToEnd();"));

    rewriteGeneratorBodyWithVars(
        "l1: try { break l1; } catch (e) { yield; } finally {}",
        "var e;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(3, 4);",
            "  return GEN_CONTEXT$0.jumpThroughFinallyBlocks(0);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 3) {",
            "  GEN_CONTEXT$0.enterFinallyBlock();",
            "  return GEN_CONTEXT$0.leaveFinallyBlock(0);",
            "}",
            "e = GEN_CONTEXT$0.enterCatchBlock();",
            "return GEN_CONTEXT$0.yield(void 0, 4)"));
  }

  @Test
  public void testFinally() {
    rewriteGeneratorBodyWithVars(
        "try {yield 1;} catch (e) {} finally {b();}",
        "var e;",
        lines(
            "if (GEN_CONTEXT$0.nextAddress == 1) {",
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(2, 3);",
            "  return GEN_CONTEXT$0.yield(1, 3);",
            "}",
            "if (GEN_CONTEXT$0.nextAddress != 2) {",
            "  GEN_CONTEXT$0.enterFinallyBlock();",
            "  b();",
            "  return GEN_CONTEXT$0.leaveFinallyBlock(0);",
            "}",
            "e = GEN_CONTEXT$0.enterCatchBlock();",
            "return GEN_CONTEXT$0.jumpTo(3);"));

    rewriteGeneratorSwitchBodyWithVars(
        lines(
            "try {",
            "  try {",
            "    yield 1;",
            "    throw 2;",
            "  } catch (x) {",
            "    throw yield x;",
            "  } finally {",
            "    yield 5;",
            "  }",
            "} catch (thrown) {",
            "  yield thrown;",
            "}"),
        "var x; var thrown;",
        lines(
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(2);",
            "  GEN_CONTEXT$0.setCatchFinallyBlocks(4, 5);",
            "  return GEN_CONTEXT$0.yield(1, 7);",
            "case 7:",
            "  throw 2;",
            "case 5:",
            "  GEN_CONTEXT$0.enterFinallyBlock(2);",
            "  return GEN_CONTEXT$0.yield(5, 8);",
            "case 8:",
            "  GEN_CONTEXT$0.leaveFinallyBlock(6);",
            "  break;",
            "case 4:",
            "  x = GEN_CONTEXT$0.enterCatchBlock();",
            "  return GEN_CONTEXT$0.yield(x, 9);",
            "case 9:",
            "  throw GEN_CONTEXT$0.yieldResult;",
            "  GEN_CONTEXT$0.jumpTo(5);",
            "  break;",
            "case 6:",
            "  GEN_CONTEXT$0.leaveTryBlock(0);",
            "  break;",
            "case 2:",
            "  thrown = GEN_CONTEXT$0.enterCatchBlock();",
            "  return GEN_CONTEXT$0.yield(thrown,0)"));
  }

  /** Tests correctness of type information after transpilation */
  @Test
  public void testYield_withTypes() {
    Node returnNode =
        testAndReturnBodyForNumericGenerator(
            "yield 1 + 2;", "", "return GEN_CONTEXT$0.yield(1 + 2, 0);");

    checkState(returnNode.isReturn(), returnNode);
    Node callNode = returnNode.getFirstChild();
    checkState(callNode.isCall(), callNode);
    // TODO(b/142881197): we can't give a more accurate type right now.
    assertNode(callNode).hasColorThat().isEqualTo(StandardColors.UNKNOWN);

    Node yieldFn = callNode.getFirstChild();
    Node jscompGeneratorContext = yieldFn.getFirstChild();
    Color generatorContext =
        Iterables.getOnlyElement(
            CodeSubTree.findFirstNode(
                    getLastCompiler().getExternsRoot(),
                    (n) -> n.matchesQualifiedName("$jscomp.generator.Context"))
                .getColor()
                .getInstanceColors());
    assertThat(jscompGeneratorContext.getColor()).isEqualTo(generatorContext);

    // Check types on "1 + 2" are still present after transpilation
    Node yieldedValue = callNode.getSecondChild(); // 1 + 2

    checkState(yieldedValue.isAdd(), yieldedValue);
    assertNode(yieldedValue).hasColorThat().isEqualTo(StandardColors.NUMBER);
    assertNode(yieldedValue.getFirstChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 1
    assertNode(yieldedValue.getSecondChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 2

    Node zero = yieldedValue.getNext();
    checkState(0 == zero.getDouble(), zero);
    assertNode(zero).hasColorThat().isEqualTo(StandardColors.NUMBER);
  }

  @Test
  public void testYieldAll_withTypes() {
    Node returnNode =
        testAndReturnBodyForNumericGenerator(
            "yield * [1, 2];", "", "return GEN_CONTEXT$0.yieldAll([1, 2], 0);");

    checkState(returnNode.isReturn(), returnNode);
    Node callNode = returnNode.getFirstChild();
    checkState(callNode.isCall(), callNode);
    // TODO(b/142881197): we can't give a more accurate type right now.
    assertNode(callNode).hasColorThat().isEqualTo(StandardColors.UNKNOWN);

    Node yieldAllFn = callNode.getFirstChild();
    checkState(yieldAllFn.isGetProp());

    // Check that the original types on "[1, 2]" are still present after transpilation
    Node yieldedValue = callNode.getSecondChild(); // [1, 2]

    checkState(yieldedValue.isArrayLit(), yieldedValue);
    assertNode(yieldedValue)
        .hasColorThat()
        .isEqualTo(getLastCompiler().getColorRegistry().get(StandardColors.ARRAY_ID)); // [1, 2]
    assertNode(yieldedValue.getFirstChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 1
    assertNode(yieldedValue.getSecondChild()).hasColorThat().isEqualTo(StandardColors.NUMBER); // 2

    Node zero = yieldedValue.getNext();
    checkState(0 == zero.getDouble(), zero);
    assertNode(zero).hasColorThat().isEqualTo(StandardColors.NUMBER);
  }

  @Test
  public void testGeneratorForIn_withTypes() {
    Node case1Node =
        testAndReturnBodyForNumericGenerator(
            "for (var i in []) { yield 3; };",
            lines(
                "var i;", //
                "var GEN_FORIN$0$0;"),
            lines(
                "if (GEN_CONTEXT$0.nextAddress == 1) {",
                "  GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn([]);",
                "}",
                "if (GEN_CONTEXT$0.nextAddress != 4) {",
                "  if (!((i = GEN_FORIN$0$0.getNext()) != null)) {",
                "    return GEN_CONTEXT$0.jumpTo(4);",
                "  }",
                "  return GEN_CONTEXT$0.yield(3, 2);",
                "}",
                ";",
                "GEN_CONTEXT$0.jumpToEnd();"));

    // GEN_FORIN$0$0 = GEN_CONTEXT$0.forIn([]);
    Node assign = case1Node.getSecondChild().getFirstFirstChild();
    checkState(assign.isAssign(), assign);

    Color propertyIterator =
        Iterables.getOnlyElement(
            CodeSubTree.findFirstNode(
                    getLastCompiler().getExternsRoot(),
                    (n) -> n.matchesQualifiedName("$jscomp.generator.Context.PropertyIterator"))
                .getColor()
                .getInstanceColors());
    assertThat(assign.getColor()).isEqualTo(propertyIterator);
    assertThat(assign.getFirstChild().getColor()).isEqualTo(propertyIterator);
    assertThat(assign.getSecondChild().getColor()).isEqualTo(propertyIterator);

    // if (!((i = GEN_FORIN$0$0.getNext()) != null)) {
    Node case2Node = case1Node.getNext();
    Node ifNode = case2Node.getSecondChild().getFirstChild();
    checkState(ifNode.isIf(), ifNode);
    Node ifCond = ifNode.getFirstChild();
    checkState(ifCond.isNot(), ifCond);
    assertNode(ifCond).hasColorThat().isEqualTo(StandardColors.BOOLEAN);
    Node ne = ifCond.getFirstChild();
    assertNode(ifCond).hasColorThat().isEqualTo(StandardColors.BOOLEAN);

    Node lhs = ne.getFirstChild(); // i = GEN_FORIN$0$0.getNext()
    assertNode(lhs)
        .hasColorThat()
        .hasAlternates(StandardColors.NULL_OR_VOID, StandardColors.STRING);
    assertNode(lhs)
        .hasColorThat()
        .hasAlternates(StandardColors.NULL_OR_VOID, StandardColors.STRING);
    assertNode(lhs)
        .hasColorThat()
        .hasAlternates(StandardColors.NULL_OR_VOID, StandardColors.STRING);
    Node getNextFn = lhs.getSecondChild().getFirstChild();
    assertNode(getNextFn).hasColorThat().isEqualTo(StandardColors.TOP_OBJECT);

    Node rhs = ne.getSecondChild();
    checkState(rhs.isNull(), rhs);
    assertNode(rhs).hasColorThat().isEqualTo(StandardColors.NULL_OR_VOID);

    // GEN_CONTEXT$0.jumpToEnd()
    Node case4Node = case2Node.getNext();
    Node jumpToEndCall = case4Node.getNext().getFirstChild();
    checkState(jumpToEndCall.isCall());

    Node jumpToEndFn = jumpToEndCall.getFirstChild();
    Node jscompGeneratorContext = jumpToEndFn.getFirstChild();

    // assertThat(jumpToEndCall.getJSType().toString()).isEqualTo("undefined");

    Color generatorContext =
        Iterables.getOnlyElement(
            CodeSubTree.findFirstNode(
                    getLastCompiler().getExternsRoot(),
                    (n) -> n.matchesQualifiedName("$jscomp.generator.Context"))
                .getColor()
                .getInstanceColors());
    assertThat(jscompGeneratorContext.getColor()).isEqualTo(generatorContext);
  }

  @Test
  public void testGeneratorTryCatch_withTypes() {
    Node case1Node =
        testAndReturnBodyForNumericGenerator(
            "try {yield 1;} catch (e) {}",
            "var e;",
            lines(
                "if (GEN_CONTEXT$0.nextAddress == 1) {",
                "  GEN_CONTEXT$0.setCatchFinallyBlocks(2);",
                "  return GEN_CONTEXT$0.yield(1, 4);",
                "}",
                "if (GEN_CONTEXT$0.nextAddress != 2) {",
                "  return GEN_CONTEXT$0.leaveTryBlock(0)",
                "}",
                "e = GEN_CONTEXT$0.enterCatchBlock();",
                "GEN_CONTEXT$0.jumpToEnd();"));
    Node case2Node = case1Node.getNext().getNext();

    // Test that "e = GEN_CONTEXT$0.enterCatchBlock();" has the unknown type
    Node eAssign = case2Node.getFirstChild();
    checkState(eAssign.isAssign(), eAssign);
    assertNode(eAssign).hasColorThat().isEqualTo(StandardColors.UNKNOWN);
    assertNode(eAssign.getFirstChild()).hasColorThat().isEqualTo(StandardColors.UNKNOWN);
    assertNode(eAssign.getSecondChild()).hasColorThat().isEqualTo(StandardColors.UNKNOWN);

    Node enterCatchBlockFn = eAssign.getSecondChild().getFirstChild();
    checkState(enterCatchBlockFn.isGetProp());
    // this is loosely typed, but it's okay because optimizations don't care about callee types.
    assertNode(enterCatchBlockFn).hasColorThat().isEqualTo(StandardColors.UNKNOWN);
  }

  @Test
  public void testGeneratorMultipleVars_withTypes() {
    Node firstStatementExprResult =
        testAndReturnBodyForNumericGenerator(
            "var x = 1, y = '2';", //
            lines(
                "var x;", //
                "var y;"),
            lines(
                "x = 1;", //
                "y = '2';",
                "GEN_CONTEXT$0.jumpToEnd();"));
    // x = 1
    final NodeSubject firstStatementSubject = assertNode(firstStatementExprResult).isExprResult();
    final NodeSubject assignASubject = firstStatementSubject.hasOneChildThat().isAssign();
    assignASubject.hasColorThat().isEqualTo(StandardColors.NUMBER);
    assignASubject.hasFirstChildThat().hasColorThat().isEqualTo(StandardColors.NUMBER);
    assignASubject.hasSecondChildThat().hasColorThat().isEqualTo(StandardColors.NUMBER);

    // y = '2';
    final NodeSubject secondStatementSubject =
        assertNode(firstStatementExprResult.getNext()).isExprResult();
    final NodeSubject assignBSubject = secondStatementSubject.hasOneChildThat().isAssign();
    assignBSubject.hasColorThat().isEqualTo(StandardColors.STRING);
    assignBSubject.hasFirstChildThat().hasColorThat().isEqualTo(StandardColors.STRING);
    assignBSubject.hasSecondChildThat().hasColorThat().isEqualTo(StandardColors.STRING);
  }

  /**
   * Tests that the given generator transpiles to the given body, and does some basic checks on the
   * transpiled generator.
   *
   * @return The first case statement in the switch inside the transpiled generator
   */
  private Node testAndReturnBodyForNumericGenerator(
      String beforeBody, String varDecls, String afterBody) {
    rewriteGeneratorBodyWithVarsAndReturnType(
        beforeBody, varDecls, afterBody, "!Generator<number>");

    Node transpiledGenerator = getLastCompiler().getJsRoot().getLastChild().getLastChild();
    Node program = getAndCheckGeneratorProgram(transpiledGenerator);

    Node programBlock = NodeUtil.getFunctionBody(program);
    return programBlock.getFirstChild();
  }

  /** Get the "program" function from a tranpsiled generator */
  private Node getAndCheckGeneratorProgram(Node genFunction) {
    Node returnNode = genFunction.getLastChild().getLastChild();
    Node callNode = returnNode.getFirstChild();
    checkState(callNode.isCall(), callNode);

    Node createGenerator = callNode.getFirstChild();
    return createGenerator.getNext().getNext();
  }
}
